import time
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import DashScopeEmbeddings
from mcp.server.fastmcp import FastMCP
from dotenv import load_dotenv
from llm_provider import get_llm

# ========================
# 日志配置（推荐：输出到文件 + stderr）
# ========================
import os
import logging
from datetime import datetime

# 创建日志目录
os.makedirs("logs", exist_ok=True)
log_filename = f"logs/rag_server_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[
        logging.FileHandler(log_filename, encoding="utf-8"),
    ]
)
logger = logging.getLogger(__name__)

# 初始化环境
load_dotenv(override=True)
dashscope_api_key = os.getenv("DASHSCOPE_API_KEY")
if not dashscope_api_key:
    raise ValueError("DASHSCOPE_API_KEY 未设置")

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

# 初始化嵌入模型
embeddings = DashScopeEmbeddings(
    model="text-embedding-v1",
    dashscope_api_key=dashscope_api_key
)

# 初始化 LLM
llm = get_llm()

# 向量库路径
VECTOR_DB_PATH = "faiss_db"

# 包装 retriever 以记录耗时
def instrument_retriever(retriever):
    async def async_invoke_with_log(query):
        start = time.time()
        docs = await retriever.ainvoke(query)
        duration = time.time() - start
        logger.info(f"检索到 {len(docs)} 个文档片段, 耗时: {duration:.2f}s")
        return docs
    return RunnableLambda(async_invoke_with_log)

# logger.info(f"正在加载向量库: {VECTOR_DB_PATH}")
if not os.path.exists(VECTOR_DB_PATH) or not os.path.exists(f"{VECTOR_DB_PATH}/index.faiss"):
    raise FileNotFoundError(f"向量数据库未找到路径: {VECTOR_DB_PATH}")
faiss_db = FAISS.load_local(
            VECTOR_DB_PATH,
            embeddings=embeddings,
            allow_dangerous_deserialization=True
        )
retriever = instrument_retriever(faiss_db.as_retriever(search_kwargs={"k": 1}))

# 构建提示词
prompt = ChatPromptTemplate.from_messages([
    ("system", """
你是一个AI助手，只能根据提供的文档内容回答问题。
禁止编造或使用外部知识。
如果上下文中没有答案，请回答：'根据提供的信息无法回答该问题。'

请尽量用清晰、分点的方式回答。

上下文：
{context}
    """),
    ("human", "{input}")
])


# ========================
# 并行检索 + LLM推理
# ========================
async def get_answer(question: str) -> str:
    """使用 RAG Chain 回答问题"""
    try:
        start = time.time()
        # 1. 检索上下文
        context_text = await retriever.ainvoke(question.strip())
        final_prompt = prompt.format(context=context_text, input=question)

        # 2. LLM 普通生成
        response = await llm.ainvoke(final_prompt)
        logger.info(f"🟢 LLM 生成完成，耗时: {time.time() - start:.2f}s")
        logger.info(f"✅ 总耗时: {time.time() - start:.2f}s")
        return StrOutputParser().invoke(response)  # 只返回字符串
    except Exception as e:
        logger.exception("❌ 问答过程中发生异常")
        return f"问答过程中出错：{str(e)}"

mcp = FastMCP("TravelPolicyRAGServer")

@mcp.tool()
async def query_travel_policy(question: str) -> str:
    """查询差旅政策文档"""
    if not question.strip():
        return "问题不能为空。"
    return await get_answer(question.strip())

if __name__ == "__main__":
    try:
        mcp.run(transport="stdio")
    except Exception as e:
        logger.critical(f"❌ 服务启动失败: {e}")
        raise